from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.edge.service import Service
from selenium.common.exceptions import (
    StaleElementReferenceException,
    TimeoutException,
    NoSuchElementException
)
import time
import os
import subprocess
import random
import winreg
import requests
import zipfile
import shutil
import platform
import pyautogui
pyautogui.PAUSE = 0.1
pyautogui.FAILSAFE = True

# ==========================================================
# 选择器配置区（页面改版时只需修改以下三处常量即可）
# ==========================================================
CONTAINER_CLASS   = "i5U4dMnB"   # 用户容器
USERNAME_CLASS    = "arnSiSbK"   # 用户名
FOLLOW_BTN_CLASS  = "zPZJ3j40"   # 关注/已关注按钮

# ===== 自动下载 EdgeDriver 的函数 =====
def get_edge_version():
    try:
        key = winreg.OpenKey(winreg.HKEY_CURRENT_USER, r"Software\Microsoft\Edge\BLBeacon")
        version, _ = winreg.QueryValueEx(key, "version")
        winreg.CloseKey(key)
        return version
    except Exception:
        try:
            key = winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, r"SOFTWARE\Microsoft\EdgeUpdate")
            version, _ = winreg.QueryValueEx(key, "version")
            winreg.CloseKey(key)
            return version
        except Exception:
            return None

def download_edgedriver(version):
    arch = "win64" if platform.machine().endswith('64') else "win32"
    url = f"https://msedgedriver.azureedge.net/{version}/edgedriver_{arch}.zip"
    driver_dir = os.path.join(os.path.expanduser("~"), "auto_edgedriver")
    os.makedirs(driver_dir, exist_ok=True)
    print(f"正在下载 EdgeDriver {version} for {arch}...")
    try:
        response = requests.get(url, stream=True, timeout=30)
        if response.status_code == 200:
            zip_path = os.path.join(driver_dir, "edgedriver.zip")
            with open(zip_path, "wb") as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
            with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                zip_ref.extractall(driver_dir)
            os.remove(zip_path)
            driver_path = os.path.join(driver_dir, "msedgedriver.exe")
            print(f"EdgeDriver 已下载到: {driver_path}")
            return driver_path
    except Exception as e:
        print(f"下载 EdgeDriver 时出错: {e}")
    return None

def auto_get_edgedriver_path():
    edge_version = get_edge_version()
    if not edge_version:
        raise Exception("无法获取 Edge 浏览器版本")
    driver_dir = os.path.join(os.path.expanduser("~"), "auto_edgedriver")
    driver_path = os.path.join(driver_dir, "msedgedriver.exe")
    if not os.path.exists(driver_path):
        downloaded_path = download_edgedriver(edge_version)
        if downloaded_path:
            return downloaded_path
        else:
            raise Exception("EdgeDriver 下载失败")
    return driver_path

# ===== 主程序开始 =====
user_data_dir = os.path.expanduser(r"~\AppData\Local\Microsoft\Edge\User Data")
profile_dir = "Default"

try:
    edge_driver_path = auto_get_edgedriver_path()
    print(f"使用自动下载的 EdgeDriver: {edge_driver_path}")
except Exception as e:
    print(f"自动下载失败: {e}")
    edge_driver_path = r"C:\Users\Devin\.cache\selenium\msedgedriver\win64\136.0.3240.47\msedgedriver.exe"
    print(f"使用备选路径: {edge_driver_path}")

def clean_browser_resources():
    try:
        subprocess.run(["taskkill", "/F", "/IM", "msedge.exe"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        print("已清理浏览器资源")
    except Exception as e:
        print(f"清理浏览器资源时出错: {e}")

clean_browser_resources()

options = webdriver.EdgeOptions()
options.add_argument(f"--user-data-dir={user_data_dir}")
options.add_argument(f"--profile-directory={profile_dir}")
service = Service(executable_path=edge_driver_path)
driver = webdriver.Edge(service=service, options=options)

driver.get("https://www.douyin.com/user/self")
original_handle = driver.current_window_handle
try:
    follow_button = WebDriverWait(driver, 1.5).until(EC.element_to_be_clickable((By.CSS_SELECTOR, 'div.vZ5A25uh.U2e7YzIW > div.UJ_bYOlF')))
    follow_button.click()
except (TimeoutException, NoSuchElementException):
    print("关注按钮选择器过期，启用备选方案：最大化窗口并点击")
    driver.maximize_window()
    pyautogui.click(610, 350)
original_handle = driver.current_window_handle
class UserScrollIterator:
    def __init__(self, driver):
        self.driver = driver
        self.last_user = None
        self.pause = 1.1
        self.user_count = 0
        self.follow_count = 0
        self.earliest_user = self.fetch_earliest_user()

    def fetch_earliest_user(self):
        try:
            dropdown = WebDriverWait(self.driver, 15).until(EC.element_to_be_clickable((By.CLASS_NAME, 'Djehbda2')))
            dropdown.click()
            earliest_btn = WebDriverWait(self.driver, 15).until(EC.element_to_be_clickable((By.ID, 'earliest')))
            earliest_btn.click()
        except (TimeoutException, NoSuchElementException):
            print("未能定位下拉菜单元素，启用备选方案")
            driver.maximize_window()
            driver.switch_to.window(original_handle)
            pyautogui.moveTo(1555, 500, duration=0.01)
            pyautogui.click(1555, 666, duration=0.01)
        time.sleep(1.5)
        user_containers = WebDriverWait(self.driver, 15).until(
            EC.presence_of_all_elements_located((By.CLASS_NAME, CONTAINER_CLASS))
        )
        earliest_user = user_containers[0].find_element(By.CLASS_NAME, USERNAME_CLASS).text
        print(f"最早关注用户名称为: {earliest_user}")
        return earliest_user

    def __iter__(self):
        try:
            time.sleep(1.5)
            dropdown = WebDriverWait(self.driver, 1.5).until(EC.element_to_be_clickable((By.CLASS_NAME, 'Djehbda2')))
            dropdown.click()
            latest_btn = WebDriverWait(self.driver, 15).until(EC.element_to_be_clickable((By.ID, 'latest')))
            latest_btn.click()
            follow_count_el = WebDriverWait(self.driver, 15).until(EC.presence_of_element_located((By.ID, 'semiTab0')))
            self.follow_count = int(''.join(filter(str.isdigit, follow_count_el.text)))
        except (TimeoutException, NoSuchElementException):
            driver.maximize_window()
            driver.switch_to.window(original_handle)
            pyautogui.moveTo(1555, 500, duration=0.5)
            pyautogui.click(1555, 600, duration=0.01)
            print("未能读取关注总数，请手动输入")
            self.follow_count = int(input("请输入关注总数: "))
        return self

    def __next__(self):
        current_end_user = None
        while current_end_user != self.earliest_user:
            try:
                time.sleep(self.pause)
                user_containers = WebDriverWait(self.driver, 15).until(
                    EC.presence_of_all_elements_located((By.CLASS_NAME, CONTAINER_CLASS))
                )
                end_container = user_containers[-1]
                self.user_count = len(user_containers)
                end_user = end_container.find_element(By.CLASS_NAME, USERNAME_CLASS).text
                current_end_user = end_user
                if current_end_user == self.last_user:
                    self.pause = min(self.pause + 0.1, 30)
                    continue
                self.driver.execute_script(
                    "arguments[0].scrollIntoView({block: 'start', behavior: 'smooth'});",
                    end_container
                )
                print(f"等待间隔: {self.pause:.1f}s")
                self.pause = max(self.pause - 0.01, 0.5)
                self.last_user = current_end_user
                return current_end_user
            except Exception as e:
                print("切换句柄会无法加载新元素，尝试切回窗口以免影响滚动操作")
                driver.switch_to.window(original_handle)
                self.pause = min(self.pause * 1.5, 255)
                time.sleep(self.pause)
                continue
        raise StopIteration("已到达最早关注用户")

def random_delay(min_delay=0.01, max_delay=0.09):
    delay = random.uniform(min_delay, max_delay)
    print(f"随机延迟 {delay:.1f}s")
    time.sleep(delay)

fail_users = {}
try:
    scroll_iter = UserScrollIterator(driver)
    earliest_user = scroll_iter.earliest_user
    iter(scroll_iter)
    follow_count = scroll_iter.follow_count
    print(f"关注总数为: {follow_count}")
    loop_count = 1
    while True:
        try:
            print(f"当前第 {loop_count} 次滚动")
            current_user = next(scroll_iter)
            print(f"滚动进度：{scroll_iter.user_count}/{scroll_iter.follow_count}，本次底部用户: {current_user}")
            loop_count += 1
        except StopIteration as e:
            print(f"滚动完成: {e}")
            break
        except Exception as e:
            print(f"发生错误: {e}")
            break

    user_containers = WebDriverWait(driver, 15).until(
        EC.presence_of_all_elements_located((By.CLASS_NAME, CONTAINER_CLASS))
    )
    processed_count = 0
    for index, container in enumerate(user_containers, start=1):
        follow_btn = container.find_element(By.CLASS_NAME, FOLLOW_BTN_CLASS)
        button_text = follow_btn.text.strip()
        while button_text == "已关注":
            driver.execute_script("arguments[0].scrollIntoView({block: 'nearest', behavior: 'smooth'});", container)
            username = container.find_element(By.CLASS_NAME, USERNAME_CLASS).text
            follow_btn.click()
            print(f"读取进度: {index}/{len(user_containers)}，用户名: {username}，状态: {button_text}，已取关用户数: {processed_count}")
            random_delay()
            follow_btn = container.find_element(By.CLASS_NAME, FOLLOW_BTN_CLASS)
            button_text = follow_btn.text.strip()
            if button_text == "关注":
                print(f"✅再次确认 {username} 取关成功")
                processed_count += 1
                break
            else:
                fail_users[username] = button_text
                print(f"❌取关失败: {username}，状态: {button_text}")
                break
    if fail_users:
        print("\n取关失败的用户:")
        for username, state in fail_users.items():
            print(f"{username}: {state}")
    print(f"\n关注数(页面显示): {follow_count}")
    print(f"关注数(实际获取): {len(user_containers)}")
    if len(user_containers) == follow_count:
        print("关注数量与列表的总用户数一致！")
    else:
        print("警告：关注数量与列表的总用户数不一致！")
except Exception as e:
    print(f"取关操作发生错误: {e}")
finally:
    input("按回车结束程序并关闭浏览器...")
    driver.quit()
